import cv2
import subprocess
import time

def main():
    # RTSP流地址
    rtsp_url = "rtsp://telchina:123qweASD@10.10.30.39:554//Streaming/Channels/1"

    # 创建视频捕获对象
    cap = cv2.VideoCapture(rtsp_url)

    # 检查是否成功打开视频流
    if not cap.isOpened():
        print("Error: Could not open video.")
        return

    # 流媒体服务器地址
    streaming_server_url = "rtmp://1.94.98.166:10935/live/stream_2"

    # 启动ffmpeg进程
    command = [
        'ffmpeg',
        '-y',
        '-f', 'rawvideo',  # Specify the input format
        '-pix_fmt', 'bgr24',
        '-s', '1280x720',  # Set the desired output size
        '-r', '30',
        '-i', '-',  # Input from stdin
        '-c:v', 'libx264',
        '-pix_fmt', 'yuv420p',
        '-preset', 'veryfast',
        '-maxrate', '2000k',
        '-bufsize', '4000k',
        '-g', '50',
        '-f', 'flv',
        streaming_server_url
    ]

    process = subprocess.Popen(command, stdin=subprocess.PIPE)

    try:
        # 循环读取视频流中的帧
        while True:
            # 读取一帧视频
            ret, frame = cap.read()

            # 如果正确读取帧，ret为True
            if not ret:
                print("Warning: Could not read frame. Retrying...")
                time.sleep(0.5)  # Add a small delay before retrying
                continue

            # 设置识别框的坐标和尺寸
            x, y, width, height = 100, 100, 200, 150

            # 在帧上画一个绿色的矩形框
            cv2.rectangle(frame, (x, y), (x + width, y + height), (0, 255, 0), 2)

            # 将处理后的帧转换为字节流
            _, buffer = cv2.imencode('.jpg', frame)
            process.stdin.write(buffer.tobytes())

            # 显示处理后的帧
            cv2.imshow('Frame with Bounding Box', frame)

            # 按 'q' 退出循环
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    except Exception as e:
        print(f"An error occurred: {e}")

    finally:
        # 释放视频捕获对象
        cap.release()
        # 关闭所有OpenCV窗口
        cv2.destroyAllWindows()
        # 关闭ffmpeg进程
        process.stdin.close()
        process.wait()

if __name__ == "__main__":
    main()
